use std::fs;
use std::path::{Path, PathBuf};
use walkdir::WalkDir;

use syn::{self, 
    parse_file, visit_mut::{self, VisitMut}, Item, ItemFn, ItemMod, UseTree, File, Visibility, Attribute
};
use quote::{quote, ToTokens};
use std::env;
use std::io::{self, Write};
use std::process::{self, Command, Stdio};
// Visitor to add pub to all items
struct AddPubVisitor;

impl VisitMut for AddPubVisitor {
    fn visit_item_mut(&mut self, item: &mut Item) {
        match item {
            Item::Fn(item_fn) => {
                item_fn.vis = Visibility::Public(syn::token::Pub::default());
            }
            Item::Struct(item_struct) => {
                item_struct.vis = Visibility::Public(syn::token::Pub::default());
            }
            Item::Enum(item_enum) => {
                item_enum.vis = Visibility::Public(syn::token::Pub::default());
            }
            Item::Trait(item_trait) => {
                item_trait.vis = Visibility::Public(syn::token::Pub::default());
            }
            Item::Impl(item_impl) => {
                // Impl blocks don't have visibility
            }
            Item::Mod(item_mod) => {
                item_mod.vis = Visibility::Public(syn::token::Pub::default());
            }
            // Add other item types as needed
            _ => {}
        }
        // Continue walking
        syn::visit_mut::visit_item_mut(self, item);
    }
}

/// 访问者结构体，用于为特定函数添加 #[test] 属性
struct AddTestAttributeVisitor {
    prefix: String,
}

impl VisitMut for AddTestAttributeVisitor {
    fn visit_item_fn_mut(&mut self, node: &mut syn::ItemFn) {
        let func_name = node.sig.ident.to_string();
        println!("Processing function: {}", func_name); // 调试打印

        if func_name.starts_with(&self.prefix) {
            println!("Function '{}' matches prefix '{}'", func_name, self.prefix); // 匹配成功

            // 检查是否已经包含 #[test] 属性
            let has_test_attr = node.attrs.iter().any(|attr| {
                if let Some(ident) = attr.path().get_ident() {
                    ident == "test"
                } else {
                    false
                }
            });

            if has_test_attr {
                println!("Function '{}' already has #[test] attribute", func_name); // 已有属性
            } else {
                // 添加 #[test] 属性
                let test_attr: Attribute = syn::parse_quote!(#[test]);
                node.attrs.push(test_attr);
                println!("Added #[test] attribute to function '{}'", func_name); // 添加属性成功
            }
        }

        // 继续遍历子节点
        visit_mut::visit_item_fn_mut(self, node);
    }

    // 如果有模块嵌套，需要处理
    fn visit_item_mod_mut(&mut self, node: &mut syn::ItemMod) {
        if let Some((_, items)) = &mut node.content {
            for item in items.iter_mut() {
                self.visit_item_mut(item);
            }
        }
    }
}

// Function to remove duplicate use statements
fn remove_duplicate_uses(file: &mut File) {
    let mut seen: Vec<String> = Vec::new();
    file.items.retain(|item| {
        if let Item::Use(use_item) = item {
            let use_str = item.to_token_stream().to_string();
            if seen.contains(&use_str) {
                false
            } else {
                seen.push(use_str);
                true
            }
        } else {
            true
        }
    });
}

// Function to check if null_test_function exists
fn has_null_test_function(file: &File) -> bool {
    for item in &file.items {
        if let Item::Fn(ItemFn { sig, .. }) = item {
            if sig.ident == "null_test_function" {
                return true;
            }
        }
    }
    false
}

// Function to add null_test_function
fn add_null_test_function(file: &mut File) {
    let func: Item = syn::parse2(quote! {
        pub fn null_test_function() {}
    }).unwrap();
    file.items.push(func);
}

/// 函数用于处理单个文件
fn process_file(path: &Path, is_test: bool) -> Option<File> {
    let content = fs::read_to_string(path).ok()?;
    let mut syntax = match parse_file(&content) {
        Ok(file) => file,
        Err(e) => {
            eprintln!("Failed to parse {}: {}", path.display(), e);
            return None;
        }
    };

    // 移除重复的 use 语句
    remove_duplicate_uses(&mut syntax);

    if is_test {
        let file_name = path.file_name()?.to_string_lossy();

        if file_name == "alloc_testing.rs" || file_name == "framework.rs" {
            // 为所有定义添加 pub
            let mut visitor = AddPubVisitor;
            visitor.visit_file_mut(&mut syntax);

            if file_name == "framework.rs" && !has_null_test_function(&syntax) {
                add_null_test_function(&mut syntax);
            }
        } else {
            // 检查文件是否位于 /test/ 目录中，并且文件名以 test_xx.rs 开头
            if let Some(parent) = path.parent() {
                println!("Parent directory: {}", parent.display()); // 调试打印
                // 使用 `ends_with` 需确保匹配的是目录名
                if parent.ends_with("test") && file_name.starts_with("test_") && file_name != "alloc_testing.rs" && file_name != "framework.rs" {
                    // 提取前缀，例如 "test_xx" 从 "test_xx.rs"
                    //let prefix = file_name.trim_start_matches("test_").trim_end_matches(".rs");
                    //let test_prefix = format!("test_{}", prefix);
                    let test_prefix = "test_".to_string();
                    println!("Using test_prefix: {}", test_prefix); // 调试打印

                    // 创建并应用访问者以添加 #[test] 属性
                    let mut test_visitor = AddTestAttributeVisitor { prefix: test_prefix };
                    test_visitor.visit_file_mut(&mut syntax);
                }
            }
        }
    }

    Some(syntax)
}

/// Function to convert syn::File to code string using `quote`
fn ast_to_string(file: &File) -> String {
    file.to_token_stream().to_string()
}

/// Function to format code using `rustfmt`
fn format_code_with_rustfmt(code: &str) -> io::Result<String> {
    // 启动 `rustfmt` 作为子进程
    let mut child = Command::new("rustfmt")
        .arg("--emit")
        .arg("stdout") // 让 rustfmt 将格式化后的代码输出到 stdout
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .spawn()?;

    // 将代码写入子进程的 stdin
    {
        let stdin = child.stdin.as_mut().ok_or_else(|| {
            io::Error::new(io::ErrorKind::Other, "Failed to open stdin")
        })?;
        stdin.write_all(code.as_bytes())?;
    }

    // 获取 rustfmt 的输出
    let output = child.wait_with_output()?;

    if output.status.success() {
        let formatted_code = String::from_utf8(output.stdout)
            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
        Ok(formatted_code)
    } else {
        let error_message = String::from_utf8_lossy(&output.stderr);
        Err(io::Error::new(
            io::ErrorKind::Other,
            format!("rustfmt failed: {}", error_message),
        ))
    }
}

/// Function to write the modified AST back to the file using rustfmt
fn write_file(path: &Path, file: &File) -> std::io::Result<()> {
    // 将 AST 转换为代码字符串
    let code = ast_to_string(file);

    // 使用 rustfmt 格式化代码
    let formatted_code = format_code_with_rustfmt(&code)
        .map_err(|e| io::Error::new(io::ErrorKind::Other, format!("rustfmt error: {}", e)))?;

    // 将格式化后的代码写入文件
    fs::write(path, formatted_code)
}

fn main() {
    // 收集命令行参数
    let args: Vec<String> = env::args().collect();

    // 检查是否提供了正确数量的参数
    if args.len() != 2 {
        eprintln!("用法: {} <target_directory>", args[0]);
        process::exit(1);
    }

    // 获取用户提供的目录路径
    let target_dir = &args[1];
    
    // 验证目录路径
    let base_dir = Path::new(target_dir);
    if !base_dir.exists() {
        eprintln!("错误: 目录 '{}' 不存在。", target_dir);
        process::exit(1);
    }

    if !base_dir.is_dir() {
        eprintln!("错误: '{}' 不是一个有效的目录。", target_dir);
        process::exit(1);
    }

    // 定义要遍历的子目录
    let dirs_to_traverse = vec!["src", "test"];

    for dir in dirs_to_traverse {
        let full_path = base_dir.join(dir);
        if !full_path.exists() || !full_path.is_dir() {
            eprintln!(
                "Directory '{}' does not exist or is not a directory.",
                full_path.display()
            );
            continue;
        }

        for entry in WalkDir::new(&full_path)
            .into_iter()
            .filter_map(|e| e.ok())
            .filter(|e| {
                e.file_type().is_file()
                    && e.path()
                        .extension()
                        .map(|ext| ext == "rs")
                        .unwrap_or(false)
            })
        {
            let path = entry.path();
            let is_test = dir == "test";
            println!("Processing {}", path.display());

            if let Some(mut ast) = process_file(path, is_test) {
                if is_test {
                    let file_name = path.file_name().unwrap().to_string_lossy();
                    if file_name == "framework.rs" {
                        // Additional check already handled in process_file
                    }
                }

                if let Err(e) = write_file(path, &ast) {
                    eprintln!("Failed to write to {}: {}", path.display(), e);
                }
            } else {
                println!("Skipping {}", path.display());
            }
        }
    }

    println!("Processing completed.");
}
